import torch
import torchvision
from torch.utils.data import DataLoader
from SK_ResNet import sk_resnet

# 数据集路径
root = "../dataset"
# 批加载图片数量
batch_size = 20

# 测试集
test_set = torchvision.datasets.CIFAR100(root=root, train=False, transform=torchvision.transforms.ToTensor(),
                                         download=True)
# loader
test_dataloader = DataLoader(test_set, batch_size=batch_size)
# 数据集长度
test_set_len = len(test_set)
print(test_set_len)

# 创建网络模型
sk_resnet = sk_resnet()
# 载入权重
sk_resnet.load_state_dict(torch.load("D:/Code/Python/DeepLearning_GPU/SK_ResNet/epoch10/module/sk_resnet_10.pth",
                                  map_location=torch.device("cpu"))
                       )
print("加载网络模型成功")

print("开始测试")
sk_resnet.eval()
total_accuracy = 0  # 总的测试正确数
test_step = 0  # 记录测试次数

with torch.no_grad():
    for data in test_dataloader:
        imgs, targets = data
        output = sk_resnet(imgs)

        accuracy = (output.argmax(1) == targets).sum()
        total_accuracy += accuracy

        test_step += 1
        if test_step % 100 == 0:
            print("测试次数：{} \taccuracy：{}/{}".format(test_step, total_accuracy, test_set_len))
print("测试集上的accuracy：{}%".format(100 * total_accuracy / test_set_len))
